package top.codeora.aiollama.agent;

import cn.hutool.core.util.StrUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import top.codeora.aiollama.agent.model.AgentState;


import java.util.ArrayList;
import java.util.List;

/**
 * 抽象代理类，用于管理代理状态执行流程
 */
@Slf4j
@Data
public abstract class BaseAgent {

    //核心属性
    private String name;

    //提示词
    private String systemPrompt;
    private String nextStepPrompt;


    //代理状态
    private AgentState state = AgentState.IDLE;

    //执行步骤控制
    private int currentStep = 0;
    private int maxSteps = 10;

    //LLM 大模型
    private ChatClient chatClient;

    //Memory 记忆（需要自主维护上下文）
    private List<Message> messageList = new ArrayList<>();

    // 循环检测
    private int duplicateThreshold = 2;

    //运行代理
    public String run(String userPrompt) {
        if (this.state != AgentState.IDLE) {
            throw new RuntimeException("Cannot run agent from state:" + this.state);
        }
        if (StrUtil.isBlank(userPrompt)) {
            throw new RuntimeException("Cannot run agent with empty user prompt");
        }
        //执行，更改状态
        this.state = AgentState.RUNNING;
        //记录消息上下文
        messageList.add(new UserMessage(userPrompt));
        ArrayList<String> results = new ArrayList<>();

        try {
            for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("执行步骤 {} / {}", stepNumber, maxSteps);
                //单步执行
                String stepResult = step();
                String result = "Step" + stepNumber + ":" + stepResult;
                results.add(result);

                // 检查是否陷入循环
                if (isStuck()) {
                    handleStuckState();
                    results.add("检测到可能的循环，已添加额外提示以避免重复");
                }
            }
            //检查是否超出步骤限制
            if (currentStep >= maxSteps) {
                state = AgentState.FINISHED;
                results.add("超出步骤限制");
            }
            return String.join("\n", results);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("Error running agent", e);
            return "Error running agent: " + e.getMessage();
        } finally {
            cleanup();
        }
    }


    protected void handleStuckState() {
        String stuckPrompt = "观察到重复响应。请考虑新的策略，避免重复已尝试过的无效路径。";
        this.nextStepPrompt = stuckPrompt + "\n" + (this.nextStepPrompt != null ? this.nextStepPrompt : "");
        log.warn("检测到智能体陷入循环状态。添加额外提示: {}", stuckPrompt);
    }

    //检测是否陷入循环
    protected boolean isStuck() {
        if (messageList.size() < 2) {
            return false;
        }
        //获取最后一条信息
        AssistantMessage lastAssistantMessage = null;
        for (int i = messageList.size() - 1; i >= 0; i--) {
            if (messageList.get(i) instanceof AssistantMessage) {
                lastAssistantMessage = (AssistantMessage) messageList.get(i);
                break;
            }
        }

        if (lastAssistantMessage == null || lastAssistantMessage.getText() == null
                || lastAssistantMessage.getText().isEmpty()) {
            return false;
        }

        //计算重复内容
        int repeat = 0;
        String lastContext = lastAssistantMessage.getText();
        for (int i = messageList.size() - 2; i >= 0; i--) {
            Message message = messageList.get(i);
            if (message instanceof AssistantMessage) {
                if (lastContext.equals(message.getText())) {
                    repeat++;
                    if (repeat >= this.duplicateThreshold) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

    /**
     * 定义单个步骤
     */
    public abstract String step();


    /**
     * 清理资源
     */
    protected void cleanup() {
        messageList.clear();
    }
}
